import numpy as np
import pandas as pd
from scipy.optimize import brute
from functools import partial
import time
import os
import sys
from scipy import optimize
import math

def find_ssr(params, df, name_col, action_col,difficult_col):
    """
    Compute the SSR for a specified set of the four parameters and the skill
    based on the students actions

    Parameters
    ----------
    params : tuple ---- The four BKT model parameters: L0, T, G, and S
    df : data frame
    name_col : either str or integer ---- Column name or index of student names in the data
    action_col : either str or integer  ---- Column name or index of student responses in the data

    Returns
    -------
    ssr : float
        The sum of squared residuals by fitting data with BKT model
    """
    
    print('L0: {}, T: {}, G: {}, S: {}, D:{}'.format(*params))
    l0, trans, guess, slip,diff = params
    columns = [name_col, action_col,difficult_col]
    prevName = 'StartStudentName'
    nguess = 1.0 - guess
    nslip = 1.0 - slip
    ssr = 0.0
    for idx, (name, action,difficult_col) in df[columns].iterrows():
        if(name != prevName):
            prevL = l0
            prevName = name
        likelihood_correct = prevL* (1.0 - slip) + (1.0 - prevL) * guess
        #ssr += (action - likelihood_correct) ** 2
        ssr +=(action*math.log(likelihood_correct) + (1-action)*math.log(1-likelihood_correct))
        if action:
            postL = (prevL * nslip) / (prevL * nslip + (1.0 - prevL) * guess)
        else:
            postL = prevL * slip / (prevL * slip + (1.0 - prevL) * nguess)
        prevL = postL + (1.0 - postL) * trans
    print(-ssr,df.shape[0],-ssr/(df.shape[0]))  
    return -ssr

def bruteforece_bkt(df, ssr_func, name_col, action_col, difficult_col,
                    l0_range, trans_range, guess_range, slip_range, dif_range, step_size):
    """
    Minimize SSR to search for an optimal set of the four parameters by brute 
    forece .

    Parameters
    ----------
    df : data frame
        The student problem anwsering data
    ssr_func : callable
        The function that calculates SSR and is to be minimized
    name_col : either str or integer
        Column name or index of student names in the data
    action_col : either str or integer
        Column name or index of student responses in the data
    l0_range : tuple
        A range tuple of the form ``(low, high)`` for parameter L0
    trans_range : tuple
        A range tuple of the form ``(low, high)`` for parameter T
    guess_range : tuple
        A range tuple of the form ``(low, high)`` for parameter G
    slip_range : tuple
        A range tuple of the form ``(low, high)`` for parameter S

    Returns
    -------
    xmin : ndarray
        A 1-D array containing the coordinates of the four BKT model parameters
        at which the SSR function is minimized
    """
    pranges = (slice(l0_range[0], l0_range[1], step_size),
               slice(trans_range[0], trans_range[1], step_size),
               slice(guess_range[0], guess_range[1], step_size),
               slice(slip_range[0], slip_range[1], step_size),
               slice(dif_range[0], dif_range[1], step_size))

    extra_args = {'df': df, 'name_col': name_col, 'action_col': action_col, 'difficult_col':difficult_col}
    ssr_partial = partial(find_ssr, **extra_args)
    time_start = time.time()
    res_bf = brute(ssr_partial, pranges, full_output = True,finish=optimize.fmin)
    time_end = time.time()
    print('Running time: {}'.format(time_end - time_start))
    return res_bf

if __name__ == '__main__':
    init=float(0.12)
    data = pd.read_csv('data/train.csv', sep = ',')
    skill_dxo = data[data.kc_reduce == 'eg']
    params_bf = bruteforece_bkt(skill_dxo, find_ssr, 
                                name_col = 'user_id', action_col = 'correct', difficult_col = "difficulty",
                                l0_range = (init, init+0.01),
                                trans_range = (0.01, 0.99),
                                guess_range = (0.01, 0.30),
                                slip_range = (0.01, 0.10),
                                dif_range = (0.01, 0.1),
                                step_size = 0.05)
    parameter = params_bf[0]
    value = params_bf[1]
    print('L0: {}, T: {}, G: {}, S: {}'.format(*parameter))
    print(value)